from diffusers import UNet2DConditionModel
import torch


def get_inputs():
    sample = torch.randn(1, 4, 64, 64).to("cuda")
    timestep = torch.randint(1, 1000, size=(1,)).to("cuda")
    encoder_hidden_states = torch.randn(1, 77, 2048).to("cuda")
    added_cond_kwargs = {"text_embeds": torch.randn(1, 1280).to("cuda"), "time_ids": torch.randn(1, 6).to("cuda")}
    return dict(
        sample=sample,
        timestep=timestep,
        encoder_hidden_states=encoder_hidden_states,
        added_cond_kwargs=added_cond_kwargs,
    )


repo_id = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors"
single_file_unet = UNet2DConditionModel.from_single_file(repo_id, pipeline_class_name="StableDiffusionXLPipeline").to(
    "cuda"
)
diffusers_unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet").to(
    "cuda"
)

with torch.no_grad():
    inputs = get_inputs()
    single_file_outs = single_file_unet(**inputs, return_dict=False)[0]
    diffusers_outs = diffusers_unet(**inputs, return_dict=False)[0]

    assert torch.allclose(single_file_outs, diffusers_outs, atol=1e-4, rtol=1e-4)